{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Hyperparameter Sweep Autoencoder",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "authorship_tag": "ABX9TyMAqBmg9UiIeKd2RhRJiRNE",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ayulockin/deepgenerativemodeling/blob/master/Hyperparameter_Sweep_Autoencoder.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9U4eqN-N5P45",
        "colab_type": "text"
      },
      "source": [
        "# Setups, Imports and Installations"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F2Pas_u35J5Q",
        "colab_type": "code",
        "outputId": "370a85b6-1496-454b-ba44-6989cb01f22e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "source": [
        "import tensorflow as tf\n",
        "print(tf.__version__)\n",
        "\n",
        "from tensorflow import keras\n",
        "import tensorflow.keras.backend as K"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "2.2.0-rc2\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EdNuACCZ5sgz",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%%capture\n",
        "!pip install wandb\n",
        "\n",
        "import wandb\n",
        "from wandb.keras import WandbCallback\n",
        "\n",
        "!wandb login 69f60a7711ce6b8bbae91ac6d15e45d6b1f1430e"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pcd2uWQJ5wRj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "%matplotlib inline\n",
        "\n",
        "from tqdm.notebook import tqdm"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ST4D3VC37ux8",
        "colab_type": "text"
      },
      "source": [
        "# Let's use MNIST Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A5mg4qlY7ncx",
        "colab_type": "code",
        "outputId": "fa238f3b-784a-4435-f85b-a711975cd380",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 109
        }
      },
      "source": [
        "from keras.datasets import mnist\n",
        "\n",
        "# MNIST dataset\n",
        "(train_img, train_label), (test_img, test_label) = mnist.load_data()\n",
        "\n",
        "image_size = train_img.shape[1]\n",
        "original_dim = image_size * image_size\n",
        "train_img = np.reshape(train_img, [-1, original_dim])\n",
        "test_img = np.reshape(test_img, [-1, original_dim])\n",
        "train_img = train_img.astype('float32') / 255\n",
        "test_img = test_img.astype('float32') / 255\n",
        "\n",
        "## Just dataset info\n",
        "print(\"X_train: \", train_img.shape)\n",
        "print(\"X_test: \", test_img.shape)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Using TensorFlow backend.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Downloading data from https://s3.amazonaws.com/img-datasets/mnist.npz\n",
            "11493376/11490434 [==============================] - 1s 0us/step\n",
            "X_train:  (60000, 784)\n",
            "X_test:  (10000, 784)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a9DLCOH2ZBub",
        "colab_type": "text"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E9_fN41-8MGo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Autoencoder:\n",
        "  def __init__(self, input_shape, original_dim, intermediate_dim, latent_space):\n",
        "    self.input_shape = input_shape\n",
        "    self.original_dim = original_dim\n",
        "    self.intermediate_dim = intermediate_dim\n",
        "    self.latent_space = latent_space\n",
        "    \n",
        "  def build_model(self):\n",
        "    ## Initialize encoder model\n",
        "    encoder_input, encoder_output = self.build_encoder()\n",
        "    ## Initialize decoder model\n",
        "    self.build_decoder()\n",
        "    ## Join encoder and decoder\n",
        "    decoder_output = self.decoder(encoder_output)\n",
        "    ## Build autoencoder model\n",
        "    return keras.models.Model(inputs=[encoder_input], outputs=[decoder_output])\n",
        "\n",
        "  def build_encoder(self):\n",
        "    ## ENCODER\n",
        "    encoder_input = keras.layers.Input(shape=self.input_shape)\n",
        "    x = keras.layers.Dense(self.intermediate_dim, activation='relu')(encoder_input)\n",
        "    ## This is latent space\n",
        "    encoder_output = keras.layers.Dense(self.latent_space)(x)\n",
        "    encoder_output = keras.layers.LeakyReLU(0.3)(encoder_output)\n",
        "    \n",
        "    self.encoder =  keras.models.Model(inputs=[encoder_input], outputs=[encoder_output])\n",
        "    return encoder_input, encoder_output\n",
        "\n",
        "  def build_decoder(self):\n",
        "    ## DECODER\n",
        "    decoder_input = keras.layers.Input(shape=self.latent_space)\n",
        "    x = keras.layers.Dense(self.intermediate_dim, activation='relu')(decoder_input)\n",
        "    ## This is reconstruction\n",
        "    decoder_output = keras.layers.Dense(self.original_dim, activation='sigmoid')(x)\n",
        "    \n",
        "    self.decoder = keras.models.Model(inputs=[decoder_input], outputs=[decoder_output])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mNrEr-wpZKUC",
        "colab_type": "text"
      },
      "source": [
        "# WandB and Callbacks"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3UEJlm4ip0oJ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ReconstructionLogger(tf.keras.callbacks.Callback):\n",
        "  def __init__(self, batch_size):\n",
        "    self.batch_size = batch_size\n",
        "\n",
        "    super(ReconstructionLogger, self).__init__()\n",
        "\n",
        "  def on_epoch_end(self, logs, epoch):\n",
        "    sample_images = test_img[:self.batch_size]\n",
        "    \n",
        "    images = []\n",
        "    reconstructions = []\n",
        "    \n",
        "    for i in range(32):\n",
        "        reconstruction = self.model.predict(sample_images[i].reshape((1,)+sample_images[i].shape))\n",
        "\n",
        "        images.append(sample_images[i].reshape(28,28))\n",
        "        reconstructions.append(reconstruction.reshape(28,28))\n",
        "\n",
        "    wandb.log({\"images\": [wandb.Image(image)\n",
        "                          for image in images]})\n",
        "    wandb.log({\"reconstructions\": [wandb.Image(reconstruction)\n",
        "                          for reconstruction in reconstructions]})"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TSmoIktbZO62",
        "colab_type": "text"
      },
      "source": [
        "# Sweep Train "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9kOA4UVgZs-X",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train():\n",
        "    # Initialize wandb with a sample project name\n",
        "    wandb.init(entity='ayush-thakur', project=\"keras-gan\")\n",
        "    \n",
        "    (X_train, y_train) = train_img, train_label\n",
        "    (X_test, y_test) = test_img, test_label\n",
        "\n",
        "    # Specify the hyperparameter to be tuned along with\n",
        "    # an initial value\n",
        "    configs = {\n",
        "        'latent_space': 2\n",
        "    }\n",
        "    \n",
        "    # Specify the other hyperparameters to the configuration\n",
        "    config = wandb.config\n",
        "    config.batch_size = 32\n",
        "    config.epochs = 10\n",
        "    \n",
        "    # Add the config item (layers) to wandb\n",
        "    if wandb.run:\n",
        "        wandb.config.update({k: v for k, v in configs.items() if k not in dict(wandb.config.user_items())})\n",
        "        configs = dict(wandb.config.user_items())\n",
        "    \n",
        "    # Define the model\n",
        "    ae = Autoencoder(input_shape=(784), \n",
        "                 original_dim=(784), \n",
        "                 intermediate_dim=512,\n",
        "                 latent_space=wandb.config.latent_space)\n",
        "    \n",
        "    model = ae.build_model()\n",
        "    \n",
        "    # Compile the model\n",
        "    model.compile(optimizer='adam',\n",
        "                  loss='mean_squared_error')\n",
        "    \n",
        "    # Train the model\n",
        "    _ = model.fit(X_train,\n",
        "              X_train, \n",
        "              epochs=config.epochs,\n",
        "              batch_size=config.batch_size,\n",
        "              callbacks=[WandbCallback(),\n",
        "                         ReconstructionLogger(config.batch_size)])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TI_JI2HBZs10",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "sweep_config = {\n",
        "  'method': 'grid',\n",
        "  'parameters': {\n",
        "      'latent_space': {\n",
        "          'values': [2, 10, 100]\n",
        "      }\n",
        "  }\n",
        "}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LWpCLwGqncdF",
        "colab_type": "code",
        "outputId": "6052046a-12d5-488a-b201-511a0aabe322",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 54
        }
      },
      "source": [
        "sweep_id = wandb.sweep(sweep_config, entity='ayush-thakur', project=\"keras-gan\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Create sweep with ID: 17bqwnpv\n",
            "Sweep URL: https://app.wandb.ai/ayush-thakur/keras-gan/sweeps/17bqwnpv\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "goZ3Z7A8nbpn",
        "colab_type": "code",
        "outputId": "396f7126-6e98-4387-b7e9-544d5c2add0a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "wandb.agent(sweep_id, function=train)"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "wandb: Agent Starting Run: dz273fbw with config:\n",
            "\tlatent_space: 2\n",
            "wandb: Agent Started Run: dz273fbw\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/ayush-thakur/keras-gan\" target=\"_blank\">https://app.wandb.ai/ayush-thakur/keras-gan</a><br/>\n",
              "                Sweep page: <a href=\"https://app.wandb.ai/ayush-thakur/keras-gan/sweeps/17bqwnpv\" target=\"_blank\">https://app.wandb.ai/ayush-thakur/keras-gan/sweeps/17bqwnpv</a><br/>\n",
              "Run page: <a href=\"https://app.wandb.ai/ayush-thakur/keras-gan/runs/dz273fbw\" target=\"_blank\">https://app.wandb.ai/ayush-thakur/keras-gan/runs/dz273fbw</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/10\n",
            "1875/1875 [==============================] - 5s 2ms/step - loss: 0.0518\n",
            "Epoch 2/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0449\n",
            "Epoch 3/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0434\n",
            "Epoch 4/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0423\n",
            "Epoch 5/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0416\n",
            "Epoch 6/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0410\n",
            "Epoch 7/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0406\n",
            "Epoch 8/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0402\n",
            "Epoch 9/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0399\n",
            "Epoch 10/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0396\n",
            "wandb: Agent Finished Run: dz273fbw \n",
            "\n",
            "wandb: Agent Starting Run: 5eqplxt0 with config:\n",
            "\tlatent_space: 10\n",
            "wandb: Agent Started Run: 5eqplxt0\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/ayush-thakur/keras-gan\" target=\"_blank\">https://app.wandb.ai/ayush-thakur/keras-gan</a><br/>\n",
              "                Sweep page: <a href=\"https://app.wandb.ai/ayush-thakur/keras-gan/sweeps/17bqwnpv\" target=\"_blank\">https://app.wandb.ai/ayush-thakur/keras-gan/sweeps/17bqwnpv</a><br/>\n",
              "Run page: <a href=\"https://app.wandb.ai/ayush-thakur/keras-gan/runs/5eqplxt0\" target=\"_blank\">https://app.wandb.ai/ayush-thakur/keras-gan/runs/5eqplxt0</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0280\n",
            "Epoch 2/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0186\n",
            "Epoch 3/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0167\n",
            "Epoch 4/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0158\n",
            "Epoch 5/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0152\n",
            "Epoch 6/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0147\n",
            "Epoch 7/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0144\n",
            "Epoch 8/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0141\n",
            "Epoch 9/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0139\n",
            "Epoch 10/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0137\n",
            "wandb: Agent Finished Run: 5eqplxt0 \n",
            "\n",
            "wandb: Agent Starting Run: zzvj4eds with config:\n",
            "\tlatent_space: 100\n",
            "wandb: Agent Started Run: zzvj4eds\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "                Logging results to <a href=\"https://wandb.com\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
              "                Project page: <a href=\"https://app.wandb.ai/ayush-thakur/keras-gan\" target=\"_blank\">https://app.wandb.ai/ayush-thakur/keras-gan</a><br/>\n",
              "                Sweep page: <a href=\"https://app.wandb.ai/ayush-thakur/keras-gan/sweeps/17bqwnpv\" target=\"_blank\">https://app.wandb.ai/ayush-thakur/keras-gan/sweeps/17bqwnpv</a><br/>\n",
              "Run page: <a href=\"https://app.wandb.ai/ayush-thakur/keras-gan/runs/zzvj4eds\" target=\"_blank\">https://app.wandb.ai/ayush-thakur/keras-gan/runs/zzvj4eds</a><br/>\n",
              "            "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "Epoch 1/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0137\n",
            "Epoch 2/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0048\n",
            "Epoch 3/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0038\n",
            "Epoch 4/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0033\n",
            "Epoch 5/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0030\n",
            "Epoch 6/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0028\n",
            "Epoch 7/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0027\n",
            "Epoch 8/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0026\n",
            "Epoch 9/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0025\n",
            "Epoch 10/10\n",
            "1875/1875 [==============================] - 4s 2ms/step - loss: 0.0024\n",
            "wandb: Agent Finished Run: zzvj4eds \n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}